use std::ops::Deref;
use std::pin::Pin;
use std::sync::Arc;

use arrow_flight::{FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse, Ticket};
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::error::FlightError;
use arrow_flight::flight_service_server::FlightService;
use arrow_flight::sql::{Any, Command, CommandGetCatalogs, CommandGetDbSchemas, CommandGetTables, CommandStatementQuery, ProstMessageExt, SqlInfo};
use arrow_flight::sql::server::FlightSqlService;
use datafusion::arrow::array::{ArrayRef, RecordBatch, StringBuilder};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::arrow::error::ArrowError;
use datafusion::sql::parser::DFParser;
use futures_core::Stream;
use futures_util::{StreamExt, TryStreamExt};
use log::{debug, info};
use prost::Message;
use tonic::{Request, Response, Status, Streaming};

use engine::{DB, DBRef};

use crate::server::SqlServerImpl;

macro_rules! status {
    ($desc:expr, $err:expr) => {
        Status::internal(format!("{}: {} at {}:{}", $desc, $err, file!(), line!()))
    };
}

#[derive(Clone, PartialEq, Message)]
pub struct QueryTick {
    #[prost(uint64, tag="1")]
    id: u64,
}

impl ProstMessageExt for QueryTick {
    fn type_url() -> &'static str {
        "litedb/query-ick"
    }

    fn as_any(&self) -> Any {
        Any {
            type_url: Self::type_url().into(),
            value: self.encode_to_vec().into(),
        }
    }
}

#[tonic::async_trait]
impl FlightSqlService for SqlServerImpl {
    type FlightService = SqlServerImpl;

    async fn do_handshake(&self, _request: Request<Streaming<HandshakeRequest>>)
                          -> Result<
                              Response<Pin<Box<dyn Stream<Item=Result<HandshakeResponse, Status>> + Send>>>,
                              Status>
    {
        debug!(target: "SqlServerImpl", "do_handshake");
        let resp = HandshakeResponse {
            protocol_version: 0,
            payload: "LiteDB-V1.0.0.0".into(),
        };
        let s: Pin<Box<dyn Stream<Item=Result<HandshakeResponse, Status>> + Send>> = Box::pin(futures::stream::iter(vec![Ok(resp)]));
        Ok(Response::new(s))
    }

    async fn get_flight_info_catalogs(&self, query: CommandGetCatalogs, request: Request<FlightDescriptor>) -> Result<Response<FlightInfo>, Status> {
        debug!(target: "SqlServerImpl", "get_flight_info_catalogs");
        let flight_descriptor = request.into_inner();
        let ticket = Ticket {
            ticket: "fake-ticket".into(),
        };
        let endpoint = FlightEndpoint::new().with_ticket(ticket);
        let flight_info = FlightInfo::new()
            .try_with_schema(&query.into_builder().schema())
            .map_err(|e| status!("Unable to encode schema", e))?
            .with_endpoint(endpoint)
            .with_descriptor(flight_descriptor);
        Ok(Response::new(flight_info))
    }

    async fn get_flight_info_schemas(&self, _query: CommandGetDbSchemas, _request: Request<FlightDescriptor>) -> Result<Response<FlightInfo>, Status> {
        todo!()
    }

    async fn get_flight_info_tables(&self, _query: CommandGetTables, _request: Request<FlightDescriptor>) -> Result<Response<FlightInfo>, Status> {
        todo!()
    }

    async fn do_get_catalogs(&self, query: CommandGetCatalogs, _request: Request<Ticket>) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        debug!(target: "SqlServerImpl", "do_get_catalogs");
        let mut builder = query.into_builder();
        builder.append("test");
        let schema = builder.schema();
        let batch = builder.build();
        let stream = FlightDataEncoderBuilder::new()
            .with_schema(schema)
            .build(futures::stream::once(async { batch }))
            .map_err(Status::from);
        Ok(Response::new(Box::pin(stream)))
    }

    async fn do_get_schemas(&self, _query: CommandGetDbSchemas, _request: Request<Ticket>) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        todo!()
    }

    async fn do_get_tables(&self, _query: CommandGetTables, _request: Request<Ticket>) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        todo!()
    }

    async fn register_sql_info(&self, id: i32, result: &SqlInfo) {
        debug!(target: "flight-sql", "register_sql_info id: {}, result: {:?}", id, result);
    }

    async fn get_flight_info_statement(
        &self,
        query: CommandStatementQuery,
        request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        info!(target: "flight-sql", "get_flight_info_statement query: {:?}, request: {:?}", &query, &request);
        let flight_descriptor = request.into_inner();
        debug!("execute query: {:?}", &*query.query);
        let (id, schema) = self.query(&*query.query).await.map_err(|e| Status::internal(format!(
            "query err: {:?}", e
        )))?;
        let query_tick = QueryTick {
            id,
        };
        info!("query id: {}, query_tick: {:?}", id, &query_tick.encode_to_vec());
        let ticket = Ticket {
            ticket: query_tick.as_any().encode_to_vec().into()
        };
        info!("schema: {:?}, tick: {:?}", &schema, &ticket);
        let endpoint = FlightEndpoint::new().with_ticket(ticket);
        let flight_info = FlightInfo::new()
            .try_with_schema(&schema)
            .map_err(|e| status!("schema error", e))?
            .with_endpoint(endpoint)
            .with_descriptor(flight_descriptor);
        Ok(Response::new(flight_info))
    }

    // async fn do_get_fallback(
    //     &self,
    //     _request: Request<Ticket>,
    //     message: Any,
    // ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
    //     let batch = fake_catalog().unwrap();
    //     let stream = FlightDataEncoderBuilder::new()
    //         .with_schema(batch.schema())
    //         .build(futures::stream::once(async { Ok(batch) }))
    //         .map_err(Status::from);
    //     Ok(Response::new(Box::pin(stream)))
    // }

    async fn do_get_fallback(
        &self,
        request: Request<Ticket>,
        message: Any,
    ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
        // let batch = fake_catalog().unwrap();
        debug!("do_get_fallback, req: {:?}, message: {:?}", &request, &message);
        let query_tick: QueryTick = message.unpack().unwrap().unwrap();
        debug!("query_tick: {:?}", &query_tick);
        // let tick = &*request.get_ref().ticket;
        // let id_str = std::str::from_utf8(tick).unwrap();
        // let id: u64 = id_str.parse().unwrap();
        let (schema, data) = self.query_mgr.get_data(query_tick.id).await.unwrap();
        // let s = Pin::into_inner_unchecked(data);
        let data1 = data
            .map(|item| item.map_err(|e| FlightError::from_external_error(Box::new(e))))
            .boxed();
        info!("schema: {:?}", &schema);

        let stream = FlightDataEncoderBuilder::new()
            .with_schema(schema)
            // .build(futures::stream::once(async { Ok(batch) }))
            .build(data1)
            .map_err(Status::from);
        Ok(Response::new(Box::pin(stream)))
    }
}

fn fake_catalog() -> Result<RecordBatch, ArrowError> {
    let schema = Schema::new(vec![
        Field::new("catalog", DataType::Utf8, false)
    ]);
    let mut builder = StringBuilder::new();
    builder.append_value("test");
    let cols = vec![Arc::new(builder.finish()) as ArrayRef];
    RecordBatch::try_new(Arc::new(schema), cols)
}

#[cfg(test)]
mod test {}